#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Time    : 2021/6/29 17:00
# @Author  : LiShan
# @Email   : lishan_1997@126.com
# @File    : bg_dif.py
# @Note    : this is note

import os
import time
import cv2.cv2 as cv

save_path = "./mtd"
width = 480
height = 360

path_list = [
    save_path,
    save_path + "/fram",
    save_path + "/roi",
    save_path + "/gray",
    save_path + "/fgmask",
    save_path + "/binary",
    save_path + "/open",
    save_path + "/inflation",
    save_path + "/contour",
    save_path + "/contour_choose",
    save_path + "/object",
    save_path + "/object_choose",
]

for path in path_list:
    if os.path.exists(path):
        pass
    else:
        os.mkdir(path)


class Detected:
    def __init__(self):
        self.video = "./video/test.mp4"
        self.width = width
        self.height = height

    # 目标预检测
    def filter_contours(self, contours, width=480, height=360, rw=None, rh=None, ra=None, rwh=None):
        if rw is None:
            rw = [0.05, 0.4]
        if rh is None:
            rh = [0.05, 0.4]
        if ra is None:
            ra = [0.01, 0.3]
        if rwh is None:
            rwh = [0.5, 2]
        filter_contours = []
        for cnt in contours:
            # 依次取出每一条轮廓，计算点集或灰度图像的非零像素的右上边界矩形
            x, y, w, h = cv.boundingRect(cnt)
            area = cv.contourArea(cnt)
            img_area = width * height
            w_rate = w / width
            h_rate = h / height
            wh_rate = w / h
            area_rate = area / img_area
            # 过滤掉小物体、大物体
            c1 = area_rate < ra[0] or area_rate > ra[1]
            c2 = w_rate < rw[0] or w_rate > rw[1]
            c3 = h_rate < rh[0] or h_rate > rh[1]
            c4 = wh_rate < rwh[0] or wh_rate > rwh[1]
            if c1 or c2 or c3 or c4:
                continue
            else:
                filter_contours.append(cnt)
        return filter_contours

    def pre_object_detect(self, contours):
        pre_object = []
        for cnt in contours:
            # 依次取出每一条轮廓，计算点集或灰度图像的非零像素的右上边界矩形
            x, y, w, h = cv.boundingRect(cnt)
            rect = x, y, w, h
            pre_object.append(rect)
        return pre_object

    # 判断目标框位置，删除内部框
    def is_inside(self, o, i):
        # o为第一个框，i为第二个框
        ox, oy, ow, oh = o
        ix, iy, iw, ih = i
        # 如果第一个矩形框被完全包含在第二个矩形框中，可确定第一个矩形框应该被丢弃
        return ox > ix and oy > iy and ox + ow < ix + iw and oy + oh < iy + ih

    # 遍历检测结果来丢弃不含有检测目标的区域
    def dropout(self, pre_object):
        object = []
        num = 0
        # ri、qi为框的编号，r、q为对应框的x，y，w，h坐标参数
        for ri, r in enumerate(pre_object):
            for qi, q in enumerate(pre_object):
                if ri != qi and self.is_inside(r, q):
                    break
                else:
                    num += 1
            if num == len(pre_object):
                object.append(r)
            num = 0
        return object

    # 运动目标检测
    def target_detected(self, video=None):
        if video is None:
            video = self.video
        # 获得运行环境CPU的核心数
        nthreads = cv.getNumberOfCPUs()
        # 设置线程数
        cv.setNumThreads(nthreads)
        bw_mask = cv.imread("./lane" + "/bw_mask.jpg")
        color_mask = cv.imread("./lane" + "/color_mask.jpg")
        cap = cv.VideoCapture(video)
        fps = cap.get(cv.CAP_PROP_FPS)
        # 背景差分法（混合高斯建模）创建一个对象
        subtractor = cv.createBackgroundSubtractorMOG2(detectShadows=False)
        frame_num = 0
        tick = time.time()
        try:
            while cap.isOpened():
                start = time.time()
                ret, frame = cap.read()
                if not ret:
                    pass
                else:
                    frame_num += 1
                    # 调整图片大小：原图、（宽，高）
                    frame = cv.resize(frame, (self.width, self.height))
                    cv.imwrite(save_path + "/fram/" + str(frame_num) + "_frame" + ".jpg", frame)
                    # 原图
                    origin = frame
                    # 感兴趣区域提取
                    roi = cv.bitwise_and(bw_mask, frame)
                    cv.imwrite(save_path + "/roi/" + str(frame_num) + "_roi" + ".jpg", roi)
                    # 转为灰度图
                    gray = cv.cvtColor(roi, cv.COLOR_BGR2GRAY)
                    cv.imwrite(save_path + "/gray/" + str(frame_num) + "_gray" + ".jpg", gray)
                    # 根据背景差分法提取前景腌膜
                    fgmask = subtractor.apply(gray)
                    # new_fgmask = cv.bitwise_not(fgmask)
                    cv.imwrite(save_path + "/fgmask/" + str(frame_num) + "_fgmask" + ".jpg", fgmask)
                    # 二值化
                    # ret, binary = cv.threshold(fgmask, 125, 255, cv.THRESH_BINARY)
                    # binary = cv.adaptiveThreshold(fgmask, 255, cv.ADAPTIVE_THRESH_MEAN_C, cv.THRESH_BINARY, 11, 5)
                    binary = cv.adaptiveThreshold(fgmask, 255, cv.ADAPTIVE_THRESH_GAUSSIAN_C, cv.THRESH_BINARY, 5, 3)
                    binary = cv.bitwise_not(binary)
                    cv.imwrite(save_path + "/binary/" + str(frame_num) + "_binary" + ".jpg", binary)
                    # 形态学操作
                    element = cv.getStructuringElement(cv.MORPH_ELLIPSE, (1, 1))
                    element2 = cv.getStructuringElement(cv.MORPH_ELLIPSE, (1, 1))
                    # 开运算
                    open = cv.morphologyEx(binary, cv.MORPH_OPEN, element)
                    cv.imwrite(save_path + "/open/" + str(frame_num) + "_open" + ".jpg", open)
                    # 膨胀运算
                    inflation = cv.dilate(open, element2)
                    cv.imwrite(save_path + "/inflation/" + str(frame_num) + "_inflation" + ".jpg", inflation)
                    # 轮廓预检，轮廓跟踪算法（Suzuki，1985）
                    contour_origin = origin.copy()
                    contours, hierarchy = cv.findContours(inflation, cv.RETR_TREE, cv.CHAIN_APPROX_NONE)
                    cnt = contours[0:len(contours)]
                    cv.drawContours(contour_origin, cnt, -1, (0, 0, 255), 2)
                    cv.imwrite(save_path + "/contour/" + str(frame_num) + "_contour" + ".jpg", contour_origin)
                    # 轮廓筛选
                    contour_origin_choose = origin.copy()
                    cnt = self.filter_contours(contours)
                    cv.drawContours(contour_origin_choose, cnt, -1, (0, 0, 255), 2)
                    cv.imwrite(save_path + "/contour_choose/" + str(frame_num) + "_contour_choose" + ".jpg",
                               contour_origin_choose)
                    # 目标预检
                    obj = origin.copy()
                    pre_object = self.pre_object_detect(cnt)
                    for i in range(len(pre_object)):
                        x, y, w, h = pre_object[i]
                        cv.rectangle(obj, (x, y), (x + w, y + h), (0, 255, 0), 2)
                    cv.imwrite(save_path + "/object/" + str(frame_num) + "_object" + ".jpg", obj)
                    # 筛选目标
                    obj_choose = origin.copy()
                    object = self.dropout(pre_object)
                    for i in range(len(object)):
                        x, y, w, h = object[i]
                        cv.rectangle(obj_choose, (x, y), (x + w, y + h), (0, 255, 0), 2)
                    cv.imwrite(save_path + "/object_choose/" + str(frame_num) + "_object_choose" + ".jpg", obj_choose)
                    # cv.imshow("object", frame)
                    end = time.time()
                    t = end - start
                    print("\r正在处理第%d帧, 处理速度：%d FPS, 剩余时间：%.3f秒" %
                          (frame_num, int(1/t), ((3000-frame_num)*t)), end=" ")
                    if frame_num == 3000:
                        print("处理完成, 用时%.3f秒" % (end-tick))
                        return 0
                    k = cv.waitKey(1) & 0xff
                    if k == 27:
                        break
            cap.release()
            cv.destroyAllWindows()
        except:
            pass


if __name__ == '__main__':
    dt = Detected()
    dt.target_detected("./video/background_model.mp4")
